/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.schemas.transforms;



import com.google.auto.value.AutoValue;

import com.bff.gaia.unified.sdk.annotations.Experimental;

import com.bff.gaia.unified.sdk.annotations.Experimental.Kind;

import com.bff.gaia.unified.sdk.coders.CannotProvideCoderException;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.CoderRegistry;

import com.bff.gaia.unified.sdk.coders.RowCoder;

import com.bff.gaia.unified.sdk.schemas.FieldAccessDescriptor;

import com.bff.gaia.unified.sdk.schemas.FieldTypeDescriptors;

import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.sdk.schemas.SchemaCoder;

import com.bff.gaia.unified.sdk.schemas.utils.SelectHelpers;

import com.bff.gaia.unified.sdk.transforms.Combine.CombineFn;

import com.bff.gaia.unified.sdk.transforms.CombineFns;

import com.bff.gaia.unified.sdk.transforms.CombineFns.CoCombineResult;

import com.bff.gaia.unified.sdk.transforms.CombineFns.ComposedCombineFn;

import com.bff.gaia.unified.sdk.transforms.SerializableFunction;

import com.bff.gaia.unified.sdk.transforms.SerializableFunctions;

import com.bff.gaia.unified.sdk.transforms.SimpleFunction;

import com.bff.gaia.unified.sdk.values.Row;

import com.bff.gaia.unified.sdk.values.TupleTag;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Lists;



import javax.annotation.Nullable;

import java.io.Serializable;

import java.util.List;

import java.util.stream.Collectors;



/** This is the builder used by {@link Group} to build up a composed {@link CombineFn}. */

@Experimental(Kind.SCHEMAS)

class SchemaAggregateFn {

  static <T> Inner<T> create() {

    return new AutoValue_SchemaAggregateFn_Inner.Builder<T>()

        .setFieldAggregations(Lists.newArrayList())

        .build();

  }



  /** Implementation of {@link #create}. */

  @AutoValue

  abstract static class Inner<T> extends CombineFn<T, Object[], Row> {

    // Represents an aggregation of one or more fields.

    static class FieldAggregation<FieldT, AccumT, OutputT> implements Serializable {

      FieldAccessDescriptor fieldsToAggregate;

      // The specification of the output field.

      private final Schema.Field outputField;

      // The combine function.

      private final CombineFn<FieldT, AccumT, OutputT> fn;

      // The TupleTag identifying this aggregation element in the composed combine fn.

      private final TupleTag<Object> combineTag;

      // The schema corresponding to the the subset of input fields being aggregated.

      @Nullable

	  private final Schema inputSubSchema;

      // The flattened version of inputSubSchema.

      @Nullable

	  private final Schema unnestedInputSubSchema;

      // The output schema resulting from the aggregation.

      private final Schema aggregationSchema;

      private final boolean needsUnnesting;



      FieldAggregation(

          FieldAccessDescriptor fieldsToAggregate,

          Schema.Field outputField,

          CombineFn<FieldT, AccumT, OutputT> fn,

          TupleTag<Object> combineTag) {

        this(

            fieldsToAggregate,

            outputField,

            fn,

            combineTag,

            Schema.builder().addField(outputField).build(),

            null);

      }



      FieldAggregation(

          FieldAccessDescriptor fieldsToAggregate,

          Schema.Field outputField,

          CombineFn<FieldT, AccumT, OutputT> fn,

          TupleTag<Object> combineTag,

          Schema aggregationSchema,

          @Nullable Schema inputSchema) {

        if (inputSchema != null) {

          this.fieldsToAggregate = fieldsToAggregate.resolve(inputSchema);

          this.inputSubSchema = SelectHelpers.getOutputSchema(inputSchema, this.fieldsToAggregate);

          this.unnestedInputSubSchema = Unnest.getUnnestedSchema(inputSubSchema);

          this.needsUnnesting = !inputSchema.equals(unnestedInputSubSchema);

        } else {

          this.fieldsToAggregate = fieldsToAggregate;

          this.inputSubSchema = null;

          this.unnestedInputSubSchema = null;

          this.needsUnnesting = false;

        }

        this.outputField = outputField;

        this.fn = fn;

        this.combineTag = combineTag;

        this.aggregationSchema = aggregationSchema;

      }



      // The Schema is not necessarily known when the SchemaAggregateFn is created. Once the schema

      // is known, resolve will be called with the proper schema.

      FieldAggregation<FieldT, AccumT, OutputT> resolve(Schema schema) {

        return new FieldAggregation<>(

            fieldsToAggregate, outputField, fn, combineTag, aggregationSchema, schema);

      }

    }



    abstract Builder<T> toBuilder();



    @AutoValue.Builder

    abstract static class Builder<T> {

      abstract Builder<T> setInputSchema(@Nullable Schema inputSchema);



      abstract Builder<T> setOutputSchema(@Nullable Schema outputSchema);



      abstract Builder<T> setComposedCombineFn(@Nullable ComposedCombineFn<T> composedCombineFn);



      abstract Builder<T> setFieldAggregations(List<FieldAggregation> fieldAggregations);



      abstract Inner<T> build();

    }



    abstract @Nullable

	Schema getInputSchema();



    abstract @Nullable

	Schema getOutputSchema();



    abstract @Nullable

	ComposedCombineFn<T> getComposedCombineFn();



    abstract List<FieldAggregation> getFieldAggregations();



    /** Once the schema is known, this function is called by the {@link Group} transform. */

    Inner<T> withSchema(Schema inputSchema, SerializableFunction<T, Row> toRowFunction) {

      List<FieldAggregation> fieldAggregations =

          getFieldAggregations().stream()

              .map(f -> f.resolve(inputSchema))

              .collect(Collectors.toList());



      ComposedCombineFn<T> composedCombineFn = null;

      for (int i = 0; i < fieldAggregations.size(); ++i) {

        FieldAggregation fieldAggregation = fieldAggregations.get(i);

        SimpleFunction<T, ?> extractFunction;

        Coder extractOutputCoder;

        if (fieldAggregation.unnestedInputSubSchema.getFieldCount() == 1) {

          extractFunction = new ExtractSingleFieldFunction<>(fieldAggregation, toRowFunction);

          extractOutputCoder =

              RowCoder.coderForFieldType(

                  fieldAggregation.unnestedInputSubSchema.getField(0).getType());

        } else {

          extractFunction = new ExtractFieldsFunction<>(fieldAggregation, toRowFunction);

          extractOutputCoder = RowCoder.of(fieldAggregation.inputSubSchema);

        }

        if (i == 0) {

          composedCombineFn =

              CombineFns.compose()

                  .with(

                      extractFunction,

                      extractOutputCoder,

                      fieldAggregation.fn,

                      fieldAggregation.combineTag);

        } else {

          composedCombineFn =

              composedCombineFn.with(

                  extractFunction,

                  extractOutputCoder,

                  fieldAggregation.fn,

                  fieldAggregation.combineTag);

        }

      }



      return toBuilder()

          .setInputSchema(inputSchema)

          .setComposedCombineFn(composedCombineFn)

          .setFieldAggregations(fieldAggregations)

          .build();

    }



    /** Aggregate all values of a set of fields into an output field. */

    <CombineInputT, AccumT, CombineOutputT> Inner<T> aggregateFields(

        FieldAccessDescriptor fieldsToAggregate,

        CombineFn<CombineInputT, AccumT, CombineOutputT> fn,

        String outputFieldName) {

      return aggregateFields(

          fieldsToAggregate,

          fn,

          Schema.Field.of(outputFieldName, FieldTypeDescriptors.fieldTypeForJavaType(fn.getOutputType())));

    }



    /** Aggregate all values of a set of fields into an output field. */

    <CombineInputT, AccumT, CombineOutputT> Inner<T> aggregateFields(

        FieldAccessDescriptor fieldsToAggregate,

        CombineFn<CombineInputT, AccumT, CombineOutputT> fn,

        Schema.Field outputField) {

      List<FieldAggregation> fieldAggregations = getFieldAggregations();

      TupleTag<Object> combineTag = new TupleTag<>(Integer.toString(fieldAggregations.size()));

      FieldAggregation fieldAggregation =

          new FieldAggregation<>(fieldsToAggregate, outputField, fn, combineTag);

      fieldAggregations.add(fieldAggregation);



      return toBuilder()

          .setOutputSchema(getOutputSchema(fieldAggregations))

          .setFieldAggregations(fieldAggregations)

          .build();

    }



    private Schema getOutputSchema(List<FieldAggregation> fieldAggregations) {

      Schema.Builder outputSchema = Schema.builder();

      for (FieldAggregation aggregation : fieldAggregations) {

        outputSchema.addField(aggregation.outputField);

      }

      return outputSchema.build();

    }



    /** Extract a single field from an input {@link Row}. */

    private static class ExtractSingleFieldFunction<InputT, OutputT>

        extends SimpleFunction<InputT, OutputT> {

      private final FieldAggregation fieldAggregation;

      private final SerializableFunction<InputT, Row> toRowFunction;



      private ExtractSingleFieldFunction(

          FieldAggregation fieldAggregation, SerializableFunction<InputT, Row> toRowFunction) {

        this.fieldAggregation = fieldAggregation;

        this.toRowFunction = toRowFunction;

      }



      @Override

      public OutputT apply(InputT input) {

        Row row = toRowFunction.apply(input);

        Row selected =

            SelectHelpers.selectRow(

                row,

                fieldAggregation.fieldsToAggregate,

                row.getSchema(),

                fieldAggregation.inputSubSchema);

        if (fieldAggregation.needsUnnesting) {

          selected = Unnest.unnestRow(selected, fieldAggregation.unnestedInputSubSchema);

        }

        return selected.getValue(0);

      }

    }



    /** Extract multiple fields from an input {@link Row}. */

    private static class ExtractFieldsFunction<T> extends SimpleFunction<T, Row> {

      private FieldAggregation fieldAggregation;

      private SerializableFunction<T, Row> toRowFunction;



      private ExtractFieldsFunction(

          FieldAggregation fieldAggregation, SerializableFunction<T, Row> toRowFunction) {

        this.fieldAggregation = fieldAggregation;

        this.toRowFunction = toRowFunction;

      }



      @Override

      public Row apply(T input) {

        Row row = toRowFunction.apply(input);

        return SelectHelpers.selectRow(

            row,

            fieldAggregation.fieldsToAggregate,

            row.getSchema(),

            fieldAggregation.inputSubSchema);

      }

    }



    @Override

    public Object[] createAccumulator() {

      return getComposedCombineFn().createAccumulator();

    }



    @Override

    public Object[] addInput(Object[] accumulator, T input) {

      return getComposedCombineFn().addInput(accumulator, input);

    }



    @Override

    public Object[] mergeAccumulators(Iterable<Object[]> accumulator) {

      return getComposedCombineFn().mergeAccumulators(accumulator);

    }



    @Override

    public Coder<Object[]> getAccumulatorCoder(CoderRegistry registry, Coder<T> inputCoder)

        throws CannotProvideCoderException {

      return getComposedCombineFn().getAccumulatorCoder(registry, inputCoder);

    }



    @Override

    public Coder<Row> getDefaultOutputCoder(CoderRegistry registry, Coder<T> inputCoder) {

      return SchemaCoder.of(

          getOutputSchema(), SerializableFunctions.identity(), SerializableFunctions.identity());

    }



    @Override

    public Row extractOutput(Object[] accumulator) {

      // Build a row containing a field for every aggregate that was registered.

      CoCombineResult coCombineResult = getComposedCombineFn().extractOutput(accumulator);

      Row.Builder output = Row.withSchema(getOutputSchema());

      for (FieldAggregation fieldAggregation : getFieldAggregations()) {

        Object aggregate = coCombineResult.get(fieldAggregation.combineTag);

        output.addValue(aggregate);

      }

      return output.build();

    }

  }

}